/** @file

Copyright (c) 2006-2023, Kunlun BIOS, Kunlun Technology (Beijing) Co., Ltd.. All
Rights Reserved.

You may not reproduce, distribute, publish, display, perform, modify, adapt,
transmit, broadcast, present, recite, release, license or otherwise exploit
any part of this publication in any form, by any means, without the prior
written permission of Kunlun Technology (Beijing) Co., Ltd..

Module Name:


Abstract:


Revision History:

**/

#include <Uefi.h>
#include <Exit.h>
#include <Library/FileExplorerLib.h>
#include <Library/ReportStatusCodeLib.h>
//startk-klk-alt-P000A-add//
#include <Protocol/KlHiiPopup.h>
#include <Protocol/LoadedImage.h>

//end-klk-alt-P000A-add//

STATIC EFI_CALLBACK_INFO                            *mExitCallBackInfo;

EFI_GUID   mExitFormSetGuid      = FORMSET_ID_GUID_EXIT;
UINT16             mKeyInput;
EFI_BOOT_MANAGER_LOAD_OPTION  *mBootOption;


STATIC
EFI_STATUS
EFIAPI
LocalLoadAppFunc(
  IN CONST EFI_GUID                        *NameGuid
  )
{

  EFI_STATUS                               Status;
  UINTN                                    HandleCount;
  EFI_HANDLE                               *HandleBuffer;
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH        FvFilePath;
  UINTN                                    Index;
  EFI_FIRMWARE_VOLUME2_PROTOCOL            *Fv;
  UINTN                                    Size;
  EFI_FV_FILETYPE                          Type;
  EFI_FV_FILE_ATTRIBUTES                   Attributes;
  UINT32                                   AuthenticationStatus;
  EFI_DEVICE_PATH_PROTOCOL                 *DevicePath;
  EFI_DEVICE_PATH_PROTOCOL                 *AppDevicePath;
  EFI_HANDLE                               ImageHandle;
  UINTN                                    ExitDataSize;
  CHAR16                                   *ExitData;


  HandleCount  = 0;
  HandleBuffer = NULL;

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiFirmwareVolume2ProtocolGuid,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (EFI_ERROR(Status)) {
    return Status;
  }

  EfiInitializeFwVolDevicepathNode (&FvFilePath, NameGuid);

  for (Index = 0; Index < HandleCount; Index++) {
    Status = gBS->HandleProtocol (
                    HandleBuffer[Index],
                    &gEfiFirmwareVolume2ProtocolGuid,
                    (VOID **) &Fv
                    );
    if (!EFI_ERROR (Status)) {
      Status = Fv->ReadFile (
                     Fv,
                     NameGuid,
                     NULL,
                     &Size,
                     &Type,
                     &Attributes,
                     &AuthenticationStatus
                     );
    }
    if (EFI_ERROR (Status)) {
      continue;
    }

    //
    // Create device path of Update application
    //
    DevicePath = DevicePathFromHandle (HandleBuffer[Index]);
    if (DevicePath == NULL) {
      continue;
    }

    AppDevicePath = AppendDevicePathNode (DevicePath, (EFI_DEVICE_PATH_PROTOCOL *) &FvFilePath);
    if (AppDevicePath == NULL) {
      continue;
    }

    Status = gBS->LoadImage (
                    TRUE,
                    gImageHandle,
                    AppDevicePath,
                    NULL,
                    0,
                    &ImageHandle
                    );
    FreePool (AppDevicePath);
    if (EFI_ERROR (Status)) {
      continue;
    }

    gBS->StartImage (ImageHandle, &ExitDataSize, &ExitData);

    break;
  }
  FreePool (HandleBuffer);

  return Status;
}



/**
  This function converts an input device structure to a Unicode string.

  @param DevPath                  A pointer to the device path structure.

  @return A new allocated Unicode string that represents the device path.

**/
CHAR16 *
BmDevicePathToStr (
  IN EFI_DEVICE_PATH_PROTOCOL     *DevPath
  )
{
  EFI_STATUS                       Status;
  CHAR16                           *ToText;
  EFI_DEVICE_PATH_TO_TEXT_PROTOCOL *DevPathToText;

  if (DevPath == NULL) {
    return NULL;
  }

  Status = gBS->LocateProtocol (
                  &gEfiDevicePathToTextProtocolGuid,
                  NULL,
                  (VOID **) &DevPathToText
                  );
  ASSERT_EFI_ERROR (Status);
  ToText = DevPathToText->ConvertDevicePathToText (
                            DevPath,
                            FALSE,
                            TRUE
                            );
  ASSERT (ToText != NULL);
  return ToText;
}

/**
  This function invokes Boot Manager. It then enumerate all boot options. If
  a boot option from the Boot Manager page is selected, Boot Manager will boot
  from this boot option.

  @param  HiiHandle

**/
VOID
UpdateBootManager (
  IN EFI_HII_HANDLE                         HiiHandle
  )
{
  UINTN                         Index;
  EFI_BOOT_MANAGER_LOAD_OPTION  *BootOption;
  UINTN                         BootOptionCount;
  EFI_STRING_ID                 Token;
  CHAR16                        *HelpString;
  EFI_STRING_ID                 HelpToken;
  UINT16                        *TempStr;
  //EFI_HII_HANDLE                HiiHandle;
  UINTN                         TempSize;
  VOID                          *StartOpCodeHandle;
  VOID                          *EndOpCodeHandle;
  EFI_IFR_GUID_LABEL            *StartLabel;
  EFI_IFR_GUID_LABEL            *EndLabel;
  UINT16                        DeviceType;
  BOOLEAN                       IsLegacyOption;
  BOOLEAN                       NeedEndOp;
  UINTN                         MaxLen;

  DeviceType = (UINT16) -1;

  //
  // for better user experience
  // 1. User changes HD configuration (e.g.: unplug HDD), here we have a chance to remove the HDD boot option
  // 2. User enables/disables UEFI PXE, here we have a chance to add/remove EFI Network boot option
  //
  EfiBootManagerRefreshAllBootOption ();

  //
  // BdsDxe doesn't group the legacy boot options for the same device type
  // It's UI's choice.
  //
  //GroupMultipleLegacyBootOption4SameType ();

  BootOption = EfiBootManagerGetLoadOptions (&BootOptionCount, LoadOptionTypeBoot);

  //HiiHandle = gBootManagerPrivate.HiiHandle;

  //
  // Allocate space for creation of UpdateData Buffer
  //
  StartOpCodeHandle = HiiAllocateOpCodeHandle ();
  ASSERT (StartOpCodeHandle != NULL);

  EndOpCodeHandle = HiiAllocateOpCodeHandle ();
  ASSERT (EndOpCodeHandle != NULL);

  //
  // Create Hii Extend Label OpCode as the start opcode
  //
  StartLabel = (EFI_IFR_GUID_LABEL *) HiiCreateGuidOpCode (StartOpCodeHandle, &gEfiIfrTianoGuid, NULL, sizeof (EFI_IFR_GUID_LABEL));
  StartLabel->ExtendOpCode = EFI_IFR_EXTEND_OP_LABEL;
  StartLabel->Number       = LABEL_BOOT_OPTION;

  //
  // Create Hii Extend Label OpCode as the end opcode
  //
  EndLabel = (EFI_IFR_GUID_LABEL *) HiiCreateGuidOpCode (EndOpCodeHandle, &gEfiIfrTianoGuid, NULL, sizeof (EFI_IFR_GUID_LABEL));
  EndLabel->ExtendOpCode = EFI_IFR_EXTEND_OP_LABEL;
  EndLabel->Number       = LABEL_BOOT_OPTION_END;
//modify-klk-lyang-P000A-start//
  mKeyInput = 0x4567;
//modify-klk-lyang-P000A-end//
  NeedEndOp = FALSE;
  for (Index = 0; Index < BootOptionCount; Index++) {
    //
    // At this stage we are creating a menu entry, thus the Keys are reproduceable
    //
    mKeyInput++;

    //
    // Don't display hidden boot options, but retain inactive ones.
    //
//modify-klk-lyang-P000A-start//
    if ((BootOption[Index].Attributes & LOAD_OPTION_HIDDEN) != 0 && StrCmp(BootOption[Index].Description, L"UEFI Shell")) {
      continue;
    }
//modify-klk-lyang-P000A-end//


    //
    // Group the legacy boot option in the sub title created dynamically
    //
    IsLegacyOption = (BOOLEAN) (
                       (DevicePathType (BootOption[Index].FilePath) == BBS_DEVICE_PATH) &&
                       (DevicePathSubType (BootOption[Index].FilePath) == BBS_BBS_DP)
                       );

    if (!IsLegacyOption && NeedEndOp) {
      NeedEndOp = FALSE;
      HiiCreateEndOpCode (StartOpCodeHandle);
    }

    if (IsLegacyOption && DeviceType != ((BBS_BBS_DEVICE_PATH *) BootOption[Index].FilePath)->DeviceType) {
      if (NeedEndOp) {
        HiiCreateEndOpCode (StartOpCodeHandle);
      }

      DeviceType = ((BBS_BBS_DEVICE_PATH *) BootOption[Index].FilePath)->DeviceType;
      Token      = HiiSetString (
                     HiiHandle,
                     0,
                     mDeviceTypeStr[
                     MIN (DeviceType & 0xF, ARRAY_SIZE (mDeviceTypeStr) - 1)
                     ],
                     NULL
                     );
      HiiCreateSubTitleOpCode (StartOpCodeHandle, Token, 0, 0, 1);
      NeedEndOp = TRUE;
    }

    ASSERT (BootOption[Index].Description != NULL);

    Token = HiiSetString (HiiHandle, 0, BootOption[Index].Description, NULL);

    TempStr = BmDevicePathToStr (BootOption[Index].FilePath);
    TempSize = StrSize (TempStr);
    HelpString = AllocateZeroPool (TempSize + StrSize (L"Device Path : "));
    MaxLen = (TempSize + StrSize (L"Device Path : "))/sizeof(CHAR16);
    ASSERT (HelpString != NULL);
    StrCatS (HelpString, MaxLen, L"Device Path : ");
    StrCatS (HelpString, MaxLen, TempStr);

    HelpToken = HiiSetString (HiiHandle, 0, HelpString, NULL);

    HiiCreateActionOpCode (
      StartOpCodeHandle,
      mKeyInput,
      Token,
      HelpToken,
      EFI_IFR_FLAG_CALLBACK,
      0
      );
  }

  if (NeedEndOp) {
    HiiCreateEndOpCode (StartOpCodeHandle);
  }

  HiiUpdateForm (
    HiiHandle,
    &mExitFormSetGuid,
    ROOT_FORM_ID,
    StartOpCodeHandle,
    EndOpCodeHandle
    );

  HiiFreeOpCodeHandle (StartOpCodeHandle);
  HiiFreeOpCodeHandle (EndOpCodeHandle);
  //EfiBootManagerFreeLoadOptions (BootOption, BootOptionCount);
  mBootOption = BootOption;
}

/**
  Brief description of ExitCallbackRoutine.

  @param  This
  @param  Action
  @param  QuestionId
  @param  Type
  @param  Value
  @param  ActionRequest

  @retval EFI_SUCCESS   Function successful returned.
**/
EFI_STATUS
EFIAPI
ExitCallbackRoutine (
  IN  CONST EFI_HII_CONFIG_ACCESS_PROTOCOL   *This,
  IN  EFI_BROWSER_ACTION                     Action,
  IN  EFI_QUESTION_ID                        QuestionId,
  IN  UINT8                                  Type,
  IN  EFI_IFR_TYPE_VALUE                     *Value,
  OUT EFI_BROWSER_ACTION_REQUEST             *ActionRequest
  )
{
  EFI_STATUS                            Status;
  EFI_BOOT_MANAGER_LOAD_OPTION          *BootOption;
  UINTN                                 BootOptionCount;
  EFI_INPUT_KEY                         Key;
  EFI_HII_POPUP_PROTOCOL                *PopupHandler;
  EFI_HII_POPUP_SELECTION               UserSelection;

  //if ((Action == EFI_BROWSER_ACTION_CHANGING) && (QuestionId == EXIT_UPDATE_BIOS_ID)) {
  if ((Action == EFI_BROWSER_ACTION_CHANGING) 
     && ((QuestionId == SETUP_PARA_IMPORT_ID)
         || (QuestionId == SETUP_PARA_EXPORT_ID)
         )) {
    Action = EFI_BROWSER_ACTION_CHANGED;
  }
  if (Action != EFI_BROWSER_ACTION_CHANGED) {
    return EFI_UNSUPPORTED;
  }
  *ActionRequest = EFI_BROWSER_ACTION_REQUEST_NONE;
  Status = gBS->LocateProtocol (&gEfiHiiPopupProtocolGuid, NULL, (VOID **) &PopupHandler);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  switch (QuestionId) {
    
  case SETUP_PARA_IMPORT_ID:
    LocalLoadAppFunc (PcdGetPtr (PcdSetupParaImportApp));
    if (CheckSetupHiiDataManager ()) {
      gSetupDataManagment->LayoutIsChanged  = TRUE;
      gSetupDataManagment->Firstin          = TRUE;
      gSetupDataManagment->RepaintFrameLine = TRUE;
    }
    return EFI_SUCCESS;

  case SETUP_PARA_EXPORT_ID:
    LocalLoadAppFunc (PcdGetPtr (PcdSetupParaExportApp));
    if (CheckSetupHiiDataManager ()) {
      gSetupDataManagment->LayoutIsChanged  = TRUE;
      gSetupDataManagment->Firstin          = TRUE;
      gSetupDataManagment->RepaintFrameLine = TRUE;
    }
    return EFI_SUCCESS;

  case EXIT_RESET_SYSTEM_ID:
    PopupHandler->CreatePopup (
                    PopupHandler,
                    EfiHiiPopupStyleNull,
                    EfiHiiPopupTypeYesNo,
                    mExitCallBackInfo->HiiHandle,
                    STRING_TOKEN (STR_IF_RESET_SYSTEM),
                    &UserSelection
                    );
    if (UserSelection == EfiHiiPopupSelectionYes) {
      gRT->ResetSystem (EfiResetCold, Status, 0, NULL);
    }
    return EFI_SUCCESS;
  case EXIT_SHUTDOWN_SYSTEM_ID:
    PopupHandler->CreatePopup (
                    PopupHandler,
                    EfiHiiPopupStyleNull,
                    EfiHiiPopupTypeYesNo,
                    mExitCallBackInfo->HiiHandle,
                    STRING_TOKEN (STR_IF_SHUTDOWN_SYSTEM),
                    &UserSelection
                    );
    if (UserSelection == EfiHiiPopupSelectionYes) {
      gRT->ResetSystem (EfiResetShutdown, Status, 0, NULL);
    }
    return EFI_SUCCESS;
  default:
    break;
  }

  BootOption = EfiBootManagerGetLoadOptions (&BootOptionCount, LoadOptionTypeBoot);
  if (mBootOption != NULL) {
    BootOption = mBootOption;
  } else {
    BootOption = EfiBootManagerGetLoadOptions (&BootOptionCount, LoadOptionTypeBoot);
  }

  //
  // Clear  the  screen  before.
  //
  gST->ConOut->SetAttribute (gST->ConOut, EFI_TEXT_ATTR (EFI_LIGHTGRAY, EFI_BLACK));
  gST->ConOut->ClearScreen (gST->ConOut);

  //
  //check any reset required change is applied? if yes, reset system
  //
  //BmSetupResetReminder ();

  //
  // parse the selected option
  //
  EfiBootManagerBoot (&BootOption[QuestionId - 0x4567 - 1]);

  if (EFI_ERROR (BootOption[QuestionId - 0x4567 - 1].Status)) {
    gST->ConOut->OutputString (
                  gST->ConOut,
                  HiiGetString (mExitCallBackInfo->HiiHandle, STRING_TOKEN (STR_ANY_KEY_CONTINUE), NULL)
                  );
    gST->ConIn->ReadKeyStroke (gST->ConIn, &Key);
  }

  if (CheckSetupHiiDataManager()) {
    gSetupDataManagment->LayoutIsChanged = TRUE;
    gSetupDataManagment->Firstin = TRUE;
    gSetupDataManagment->RepaintFrameLine = TRUE;
  }

  if (mBootOption == NULL) {
    EfiBootManagerFreeLoadOptions (BootOption, BootOptionCount);
  }
  gST->ConOut->EnableCursor (gST->ConOut, FALSE);

  return EFI_SUCCESS;
}


/**
  Brief description of InstallExitCallbackRoutine.

  @param  DriverHandle
  @param  HiiHandle

  @retval EFI_SUCCESS   Function successful returned.
**/
EFI_STATUS
InstallExitCallbackRoutine (
  IN EFI_HANDLE                             DriverHandle,
  IN EFI_HII_HANDLE                         HiiHandle
  )
{
  EFI_STATUS                                Status;
  EFI_GUID                                  FormsetGuid = FORMSET_ID_GUID_EXIT;

  mExitCallBackInfo = AllocatePool (sizeof (EFI_CALLBACK_INFO));
  if (mExitCallBackInfo == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  mBootOption =NULL;
  mExitCallBackInfo->Signature                    = EFI_CALLBACK_INFO_SIGNATURE;
  mExitCallBackInfo->DriverCallback.ExtractConfig = GenericExtractConfig;
  mExitCallBackInfo->DriverCallback.RouteConfig   = GenericRouteConfig;
  mExitCallBackInfo->DriverCallback.Callback      = ExitCallbackRoutine;
  mExitCallBackInfo->HiiHandle                    = HiiHandle;
  CopyGuid (&mExitCallBackInfo->FormsetGuid, &FormsetGuid);

   //
  // Install protocol interface
  //
  Status = gBS->InstallProtocolInterface (
                  &DriverHandle,
                  &gEfiHiiConfigAccessProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  &mExitCallBackInfo->DriverCallback
                  );
  ASSERT_EFI_ERROR (Status);
  UpdateBootManager(HiiHandle);
  return Status;

}
